use ndarray::prelude::*;
use std::collections::HashMap;

pub struct COO<A: Clone> {
    data: Array1<A>,
    row_indices: Array1<usize>, // assume sorted
    col_indices: Array1<usize>,
    nrows: usize,
    ncols: usize,
}

impl<A> COO<A>
where
    A: Clone + std::marker::Send + std::marker::Sync,
{
    pub fn to_ndarray(&self) -> (Array1<A>, Array1<usize>, Array1<usize>) {
        (
            self.data.clone(),
            self.row_indices.clone(),
            self.col_indices.clone(),
        )
    }

    pub fn nnz(&self) -> usize {
        self.data.len()
    }

    pub fn shape(&self) -> (usize, usize) {
        (self.nrows, self.ncols)
    }

    pub fn new(
        data: &ArrayView1<A>,
        row_indices: &ArrayView1<usize>,
        col_indices: &ArrayView1<usize>,
        nrows: usize,
        ncols: usize,
    ) -> COO<A> {
        let mut indices = (0..row_indices.len()).collect::<Vec<usize>>();
        indices.sort_by_key(|&i| &row_indices[i]);

        COO::<A> {
            data: Array1::<A>::from_iter(indices.iter().map(|&i| data[i].clone())),
            row_indices: Array1::<usize>::from_iter(
                indices.iter().map(|&i| row_indices[i].clone()),
            ),
            col_indices: Array1::<usize>::from_iter(
                indices.iter().map(|&i| col_indices[i].clone()),
            ),
            nrows: nrows,
            ncols: ncols,
        }
    }

    pub fn tocsr(&self) -> CSR<A> {
        let csr_data = self.data.clone();
        let csr_indices = self.col_indices.clone();
        let mut csr_indptr = Array1::<usize>::zeros(self.nrows + 1);

        for &i in self.row_indices.iter() {
            csr_indptr[i + 1] += 1;
        }
        for i in 0..self.nrows {
            csr_indptr[i + 1] += csr_indptr[i];
        }

        CSR::<A> {
            data: csr_data,
            col_indices: csr_indices,
            row_indptr: csr_indptr,
            nrows: self.nrows,
            ncols: self.ncols,
        }
    }

    // pub fn tocsc(&self) -> CSC<A> {
    //     let mut indices = (0..self.col_indices.len()).collect::<Vec<usize>>();
    //     indices.sort_by_key(|&i| &self.col_indices[i]);

    //     let csc_data = Array1::<A>::from_iter(indices.iter().map(|&i| self.data[i].clone()));
    //     let csc_indices =
    //         Array1::<usize>::from_iter(indices.iter().map(|&i| self.row_indices[i].clone()));
    //     let mut csc_indptr = Array1::<usize>::zeros(self.ncols + 1);

    //     for &i in indices.iter() {
    //         csc_indptr[self.col_indices[i] + 1] += 1;
    //     }
    //     for i in 0..self.nrows {
    //         csc_indptr[i + 1] += csc_indptr[i];
    //     }

    //     CSC::<A> {
    //         data: csc_data,
    //         row_indices: csc_indices,
    //         col_indptr: csc_indptr,
    //         nrows: self.nrows,
    //         ncols: self.ncols,
    //     }
    // }
}

pub struct CSR<A: Clone + std::marker::Send+ std::marker::Sync> {
    data: Array1<A>,
    col_indices: Array1<usize>,
    row_indptr: Array1<usize>,
    nrows: usize,
    ncols: usize,
}

impl<A: Clone + std::marker::Send+ std::marker::Sync> CSR<A> {
    pub fn remove_element(&mut self, src: usize, dst: usize) {
        // 边界检查
        assert!(src < self.nrows, "src exceeds matrix dimensions");
        assert!(dst < self.ncols, "dst exceeds matrix dimensions");

        let row_start = self.row_indptr[src];
        let row_end = self.row_indptr[src + 1];
        
        // 在指定行的列索引中查找目标位置
        if let Some((pos, _)) = self.col_indices
            .slice(s![row_start..row_end])
            .iter()
            .enumerate()
            .find(|&(_, &c)| c == dst)
        {
            let delete_idx = row_start + pos;
            
            // 将数据转换为Vec进行删除操作
            let mut data_vec = self.data.to_vec();
            let mut col_vec = self.col_indices.to_vec();
            
            data_vec.remove(delete_idx);
            col_vec.remove(delete_idx);
            
            // 更新索引指针数组
            let mut indptr_vec = self.row_indptr.to_vec();
            for i in (src + 1)..indptr_vec.len() {
                indptr_vec[i] -= 1;
            }

            // 转换回ndarray格式
            self.data = Array1::from(data_vec);
            self.col_indices = Array1::from(col_vec);
            self.row_indptr = Array1::from(indptr_vec);
        }
    }
}

impl<A> CSR<A>
where
    A: Clone + std::marker::Send+ std::marker::Sync,
{
    pub fn nnz(&self) -> usize {
        self.data.len()
    }

    pub fn shape(&self) -> (usize, usize) {
        (self.nrows, self.ncols)
    }

    pub fn submatrix(&self, row: &Array1<usize>, col: &Array1<usize>) -> CSR<A> {
        let mut nd = Vec::<A>::new();
        let mut nc = Vec::<usize>::new();
        let mut nr = Vec::<usize>::new();
        let col_map: HashMap<usize, usize> =
            HashMap::from_iter(col.iter().enumerate().map(|(x, &y)| (y, x)));
        let mut rsum = 0;
        nr.push(rsum);
        for &r in row {
            let (a, b) = self.row(r);
            for (&c, &ref v) in a.iter().zip(b.iter()) {
                if col_map.contains_key(&c) {
                    nd.push(v.clone());
                    nc.push(*col_map.get(&c).unwrap());
                    rsum += 1;
                }
            }
            nr.push(rsum);
        }

        CSR::<A> {
            data: Array1::<A>::from(nd),
            col_indices: Array1::<usize>::from(nc),
            row_indptr: Array1::<usize>::from(nr),
            nrows: row.len(),
            ncols: col.len(),
        }
    }

    pub fn row(&self, n: usize) -> (ArrayView1<usize>, ArrayView1<A>) {
        let s = self.row_indptr[n];
        let e = self.row_indptr[n + 1];
        (self.col_indices.slice(s![s..e]), self.data.slice(s![s..e]))
    }

    pub fn transpose(&self) -> CSR<A> {
        let mut row_counts = Array1::<usize>::zeros(self.ncols);
        let mut new_row_ptr = Array1::zeros(self.ncols + 1);
        for &col in &self.col_indices {
            row_counts[col] += 1;
        }

        for i in 0..self.ncols {
            new_row_ptr[i + 1] = new_row_ptr[i] + row_counts[i];
        }

        let mut new_values = self.data.clone();
        let mut new_col_indices = Array1::zeros(self.col_indices.len());
        let mut current_pos = new_row_ptr.clone();

        for row in 0..self.nrows {
            let row_start = self.row_indptr[row];
            let row_end = self.row_indptr[row + 1];
            for i in row_start..row_end {
                let col = self.col_indices[i];
                let val = self.data[i].clone();
                let pos = current_pos[col];

                new_values[pos] = val;
                new_col_indices[pos] = row;
                current_pos[col] += 1;
            }
        }
        CSR::<A> {
            data: new_values,
            col_indices: new_col_indices,
            row_indptr: new_row_ptr,
            nrows: self.ncols,
            ncols: self.nrows,
        }
    }

    pub fn tocoo(&self) -> COO<A>
    where
        A: Clone,
    {
        let num_nonzero = self.nnz();
        let mut coo_data = self.data.clone();
        let mut coo_row_indices = Array1::<usize>::zeros(num_nonzero);
        let mut coo_col_indices = Array1::<usize>::zeros(num_nonzero);

        let mut current_ptr = 0;
        let mut current_row = 0;

        for (_i, &indptr) in self.row_indptr.iter().enumerate().skip(1) {
            for j in current_ptr..indptr {
                coo_data[j] = self.data[j].clone();
                coo_row_indices[j] = current_row;
                coo_col_indices[j] = self.col_indices[j];
            }
            current_ptr = indptr;
            current_row += 1;
        }
        COO::<A> {
            data: coo_data,
            row_indices: coo_row_indices,
            col_indices: coo_col_indices,
            nrows: self.nrows,
            ncols: self.ncols,
        }
    }
}

impl<A> CSR<A>
where
    A: Clone + std::marker::Send+ std::marker::Sync,
{
    pub fn subtract(&self, other: &CSR<A>) -> CSR<A> {
        if self.nrows != other.nrows || self.ncols != other.ncols {
            panic!(
                "Wrong shape:({},{}) and ({},{})",
                self.nrows, self.ncols, other.nrows, other.ncols
            );
        }
        let mut result_data: Vec<A> = Vec::new();
        let mut result_col_indices: Vec<usize> = Vec::new();
        let mut result_row_indptr: Vec<usize> = vec![0; self.nrows + 1];

        for row in 0..self.nrows {
            let self_start = self.row_indptr[row];
            let self_end = self.row_indptr[row + 1];

            let other_start = other.row_indptr[row];
            let other_end = other.row_indptr[row + 1];

            let mut i = self_start;
            let mut j = other_start;

            while i < self_end && j < other_end {
                let self_col = self.col_indices[i];
                let other_col = other.col_indices[j];

                if self_col < other_col {
                    result_data.push(self.data[i].clone());
                    result_col_indices.push(self_col);
                    i += 1;
                } else if self_col > other_col {
                    // result_data.push(other.data[j].clone());
                    // result_col_indices.push(other_col);
                    j += 1;
                } else {
                    i += 1;
                    j += 1;
                }
            }

            while i < self_end {
                result_data.push(self.data[i].clone());
                result_col_indices.push(self.col_indices[i]);
                i += 1;
            }

            // while j < other_end {
            //     result_data.push(other.data[j].clone());
            //     result_col_indices.push(other.col_indices[j]);
            //     j += 1;
            // }

            result_row_indptr[row + 1] = result_data.len();
        }

        CSR {
            data: Array1::from(result_data),
            col_indices: Array1::from(result_col_indices),
            row_indptr: Array1::from(result_row_indptr),
            nrows: self.nrows,
            ncols: self.ncols,
        }
    }
}

// pub struct CSC<A: Clone> {
//     data: Array1<A>,
//     row_indices: Array1<usize>,
//     col_indptr: Array1<usize>,
//     nrows: usize,
//     ncols: usize,
// }

// impl<A> CSC<A>
// where
//     A: Clone,
// {
//     pub fn nnz(&self) -> usize {
//         self.data.len()
//     }

//     pub fn shape(&self) -> (usize, usize) {
//         (self.nrows, self.ncols)
//     }

//     pub fn tocoo(&self) -> COO<A>
//     where
//         A: Clone,
//     {
//         let num_nonzero = self.nnz();
//         let mut coo_data = self.data.clone();
//         let mut coo_row_indices = Array1::<usize>::zeros(num_nonzero);
//         let mut coo_col_indices = Array1::<usize>::zeros(num_nonzero);

//         let mut current_ptr = 0;
//         let mut current_col = 0;

//         for (_i, &indptr) in self.col_indptr.iter().enumerate().skip(1) {
//             for j in current_ptr..indptr {
//                 coo_data[j] = self.data[j].clone();
//                 coo_row_indices[j] = self.row_indices[j];
//                 coo_col_indices[j] = current_col;
//             }
//             current_ptr = indptr;
//             current_col += 1;
//         }
//         COO::<A> {
//             data: coo_data,
//             row_indices: coo_row_indices,
//             col_indices: coo_col_indices,
//             nrows: self.nrows,
//             ncols: self.ncols,
//         }
//     }
// }

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn coo_works() {
        let data = Array1::from_iter(vec![1., 9., 5., 6., 3., 8., 7., 2., 2., 4.]);
        let row_ind = Array1::from_iter(vec![0, 4, 1, 3, 2, 4, 0, 4, 1, 3]);
        let col_ind = Array1::from_iter(vec![0, 0, 1, 1, 2, 2, 3, 3, 4, 4]);
        let nrows: usize = 5;
        let ncols: usize = 5;
        let spcoo = COO::<f64>::new(&data.view(), &row_ind.view(), &col_ind.view(), nrows, ncols);
        assert_eq!(spcoo.nnz(), 10);
        assert_eq!(spcoo.shape(), (5, 5));
        assert_eq!(spcoo.row_indices, array![0, 0, 1, 1, 2, 3, 3, 4, 4, 4]);
        assert_eq!(spcoo.col_indices, array![0, 3, 1, 4, 2, 1, 4, 0, 2, 3]);
        assert_eq!(spcoo.data, array![1., 7., 5., 2., 3., 6., 4., 9., 8., 2.]);
    }

    #[test]
    fn csr_works() {
        let data = array![1., 9., 5., 6., 3., 8., 7., 2., 2., 4.];
        let row_ind = array![0, 4, 1, 3, 2, 4, 0, 4, 1, 3];
        let col_ind = array![0, 0, 1, 1, 2, 2, 3, 3, 4, 4];
        let nrows: usize = 5;
        let ncols: usize = 5;
        let spcoo = COO::<f64>::new(&data.view(), &row_ind.view(), &col_ind.view(), nrows, ncols);
        let spcsr = spcoo.tocsr();
        assert_eq!(spcsr.nnz(), 10);
        assert_eq!(spcsr.shape(), (5, 5));
        assert_eq!(spcsr.row_indptr, array![0, 2, 4, 5, 7, 10]);
        assert_eq!(spcsr.col_indices, array![0, 3, 1, 4, 2, 1, 4, 0, 2, 3]);
        let spcsr_coo = spcsr.tocoo();
        assert_eq!(spcsr_coo.data, spcoo.data);
        assert_eq!(spcsr_coo.row_indices, spcoo.row_indices);
        assert_eq!(spcsr_coo.col_indices, spcoo.col_indices);
    }

    // #[test]
    // fn csc_works() {
    //     let data = array![1., 9., 5., 6., 3., 8., 7., 2., 2., 4.];
    //     let row_ind = array![0, 4, 1, 3, 2, 4, 0, 4, 1, 3];
    //     let col_ind = array![0, 0, 1, 1, 2, 2, 3, 3, 4, 4];
    //     let nrows: usize = 5;
    //     let ncols: usize = 5;
    //     let spcoo = COO::<f64>::new(&data.view(), &row_ind.view(), &col_ind.view(), nrows, ncols);
    //     let spcsc = spcoo.tocsc();
    //     assert_eq!(spcsc.nnz(), 10);
    //     assert_eq!(spcsc.shape(), (5, 5));
    //     assert_eq!(spcsc.row_indices, array![0, 4, 1, 3, 2, 4, 0, 4, 1, 3]);
    //     assert_eq!(spcsc.col_indptr, array![0, 2, 4, 6, 8, 10]);
    //     let spcsc_coo = spcsc.tocoo();
    //     assert_eq!(spcsc_coo.data, spcoo.data);
    //     assert_eq!(spcsc_coo.row_indices, spcoo.row_indices);
    //     assert_eq!(spcsc_coo.col_indices, spcoo.col_indices);
    // }
}
